import torch
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from Demo04_model import *

# 数据加载
train_data = torchvision.datasets.CIFAR10("../../../dataSet",
                                          train=True,
                                          download=True,
                                          transform=torchvision.transforms.ToTensor())

test_data = torchvision.datasets.CIFAR10("../../../dataSet",
                                         train=False,
                                         download=True,
                                         transform=torchvision.transforms.ToTensor())

# print("xxxxx{}".format());  格式化字符串
print(len(train_data))
print(len(test_data))

train_dataLoader = DataLoader(train_data, batch_size=64, shuffle=True, drop_last=True)
test_dataLoader = DataLoader(test_data, batch_size=64, shuffle=True, drop_last=True)

# 创建模型
ah = Ah()

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 优化器创建
lr = 1e-2
optimizer = torch.optim.SGD(ah.parameters(), lr)

# 记录 训练和测试步骤
total_train_step = 0
total_test_step = 0
# 设置循环次数
epoch = 10

writer = SummaryWriter("../../logs")

for i in range(epoch):
    total_loss = 0
    train_real = 0
    ah.train()
    for data in train_dataLoader:
        imgs, tables = data
        output = ah(imgs)
        train_real += sum(output.argmax(1),tables)
        # 计算损失函数
        loss = loss_fn(output, tables)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss
        total_train_step += 1
    print("epoch:{},loss:{}".format(i, total_loss / 5))
    print("epoch:{},zheng que lv:{}".format(i, train_real / 50000))
    writer.add_scalar("tran_loss",total_loss.item()/5,i)

    # 测试开始
    ah.eval()
    with torch.no_grad():
        test_loss = 0
        test_real = 0
        for data in test_dataLoader:
            imgs, tables = data
            output = ah(imgs)
            test_loss += loss_fn(output, tables)
            test_real += sum(output.argmax(1)==tables)

        print("ce shi de test_loss:{}".format(test_loss))
        print(" ce shi zheng que lv:{}".format(test_real/10000))
        writer.add_scalar("test_loss", test_loss.item(), i)
    # bao cun mo xin
    torch.save(ah,".odel/ah_{}.pth".format(i))
    print("model is solve")